import os
import json
import numpy as np
from sklearn.model_selection import train_test_split
from tqdm import tqdm

class TechStackDatasetGenerator:
    def __init__(self, output_dir="data"):
        """
        初始化技术栈数据集生成器
        
        Args:
            output_dir: 输出目录，默认为当前目录下的data
        """
        self.output_dir = output_dir
        # 定义支持的技术栈列表
        self.tech_stacks = [
            "Java基础",
            "并发编程",
            "JVM",
            "MySQL",
            "Redis",
            "Spring框架",
            "系统设计",
            "操作系统与网络",
            "算法",
            "微服务",
            "消息队列",
            "分布式系统",
            "容器技术",
            "前端技术",
            "安全",
            "性能优化",
            "数据结构"
        ]
        
        # 每个技术栈的关键词，用于生成面试题
        self.tech_keywords = {
            "Java基础": ["Java", "集合", "泛型", "反射", "注解", "异常", "IO", "面向对象", "接口", "抽象类", "继承", "多态", "封装", "HashMap", "ArrayList", "LinkedList", "String", "equals", "hashCode"],
            "并发编程": ["线程", "并发", "锁", "synchronized", "volatile", "CAS", "原子性", "可见性", "有序性", "线程池", "ThreadLocal", "AQS", "ReentrantLock", "CountDownLatch", "CyclicBarrier", "Semaphore", "并行", "Executor", "Future", "CompletableFuture"],
            "JVM": ["JVM", "垃圾回收", "GC", "内存模型", "类加载", "字节码", "JIT", "堆", "栈", "方法区", "永久代", "元空间", "新生代", "老年代", "Eden", "Survivor", "分代收集", "CMS", "G1", "ZGC", "堆外内存"],
            "MySQL": ["MySQL", "索引", "事务", "锁", "ACID", "隔离级别", "存储引擎", "InnoDB", "MyISAM", "SQL", "执行计划", "优化器", "慢查询", "分库分表", "主从复制", "读写分离", "连接", "JOIN", "GROUP BY", "HAVING"],
            "Redis": ["Redis", "缓存", "持久化", "RDB", "AOF", "过期策略", "数据结构", "String", "Hash", "List", "Set", "Zset", "分布式锁", "Lua脚本", "事务", "管道", "哨兵", "集群", "缓存穿透", "缓存击穿", "缓存雪崩"],
            "Spring框架": ["Spring", "IoC", "DI", "AOP", "Bean", "容器", "事务", "注解", "SpringMVC", "SpringBoot", "自动配置", "条件注入", "SpringCloud", "Feign", "Ribbon", "Hystrix", "Gateway", "Config", "Eureka", "循环依赖"],
            "系统设计": ["系统设计", "架构", "高可用", "高并发", "负载均衡", "限流", "熔断", "降级", "幂等性", "CAP", "BASE", "一致性", "可用性", "分区容忍性", "服务发现", "配置中心", "API网关", "服务注册", "服务发现", "服务治理"],
            "操作系统与网络": ["操作系统", "进程", "线程", "死锁", "内存管理", "文件系统", "TCP", "UDP", "HTTP", "HTTPS", "网络模型", "IO模型", "BIO", "NIO", "AIO", "Netty", "Socket", "DNS", "CDN", "三次握手", "四次挥手"],
            "算法": ["算法", "复杂度", "排序", "查找", "动态规划", "贪心", "回溯", "分治", "广度优先", "深度优先", "二叉树", "堆", "图", "最短路径", "最小生成树", "字符串匹配", "KMP", "哈希", "布隆过滤器", "LRU"],
            "微服务": ["微服务", "服务拆分", "服务治理", "注册中心", "配置中心", "API网关", "服务发现", "负载均衡", "熔断", "限流", "降级", "链路追踪", "日志中心", "监控告警", "SpringCloud", "Dubbo", "gRPC", "Thrift", "Nacos", "Eureka"],
            "消息队列": ["消息队列", "MQ", "Kafka", "RabbitMQ", "RocketMQ", "ActiveMQ", "消息中间件", "生产者", "消费者", "Topic", "Queue", "分区", "副本", "顺序消息", "事务消息", "延迟消息", "重试机制", "死信队列", "消息积压", "消息丢失"],
            "分布式系统": ["分布式", "一致性", "CAP", "BASE", "分布式事务", "2PC", "3PC", "TCC", "SAGA", "最终一致性", "共识算法", "Paxos", "Raft", "ZAB", "分布式锁", "分布式ID", "雪花算法", "时钟同步", "分布式缓存", "分布式会话"],
            "容器技术": ["容器", "Docker", "Kubernetes", "K8s", "Pod", "Service", "Deployment", "StatefulSet", "ConfigMap", "Secret", "Volume", "PV", "PVC", "命名空间", "节点", "集群", "编排", "镜像", "容器化", "微服务"],
            "前端技术": ["前端", "HTML", "CSS", "JavaScript", "TypeScript", "Vue", "React", "Angular", "Webpack", "Babel", "ESLint", "响应式", "SPA", "SSR", "PWA", "组件化", "状态管理", "路由", "虚拟DOM", "性能优化"],
            "安全": ["安全", "认证", "授权", "加密", "解密", "哈希", "签名", "SSL", "TLS", "HTTPS", "XSS", "CSRF", "SQL注入", "RBAC", "OAuth", "JWT", "单点登录", "密码学", "安全漏洞", "渗透测试"],
            "性能优化": ["性能优化", "缓存", "异步", "并行", "负载均衡", "连接池", "线程池", "内存优化", "CPU优化", "IO优化", "网络优化", "数据库优化", "索引优化", "SQL优化", "代码优化", "算法优化", "前端优化", "后端优化", "分布式优化", "架构优化"],
            "数据结构": ["数据结构", "数组", "链表", "栈", "队列", "树", "二叉树", "平衡树", "红黑树", "B树", "B+树", "堆", "图", "哈希表", "散列", "字典", "集合", "优先队列", "跳表", "前缀树"]
        }
        
        # 确保输出目录存在
        os.makedirs(self.output_dir, exist_ok=True)
    
    def generate_interview_questions(self, num_samples=5000):
        """
        生成面试题数据集
        
        Args:
            num_samples: 样本数量，默认5000
        
        Returns:
            生成的数据集，包含questions和labels两个字段
        """
        # 面试题模板，用于生成各种类型的面试题
        templates = [
            "请解释{}的原理和使用场景。",
            "{}有哪些核心特性？请详细描述。",
            "如何优化{}的性能？",
            "{}的底层实现是什么？",
            "在项目中，你是如何使用{}的？",
            "{}与{}的区别是什么？",
            "请描述{}的工作流程。",
            "如何处理{}中的常见问题？",
            "谈谈你对{}的理解。",
            "如何设计一个高性能的{}？",
            "{}的最佳实践有哪些？",
            "如何调试{}相关的问题？",
            "在高并发场景下，{}如何保证稳定性？",
            "请介绍{}的架构设计。",
            "{}的优缺点是什么？"
        ]
        
        questions = []  # 存储生成的面试题
        labels = []     # 存储对应的技术栈标签
        
        # 使用tqdm显示进度条
        for _ in tqdm(range(num_samples)):
            # 随机选择1-3个技术栈
            num_techs = np.random.randint(1, 4)
            selected_techs = np.random.choice(self.tech_stacks, size=num_techs, replace=False)
            
            # 为每个技术栈选择关键词
            keywords = []
            for tech in selected_techs:
                if tech in self.tech_keywords:
                    num_keywords = np.random.randint(1, 4)
                    tech_keywords = np.random.choice(self.tech_keywords[tech], size=min(num_keywords, len(self.tech_keywords[tech])), replace=False)
                    keywords.extend(tech_keywords)
            
            # 随机选择模板
            template = np.random.choice(templates)
            
            # 如果模板中有两个占位符，则选择两个关键词
            if template.count("{}") == 2:
                if len(keywords) >= 2:
                    kw1, kw2 = np.random.choice(keywords, size=2, replace=False)
                    question = template.format(kw1, kw2)
                else:
                    # 如果关键词不够，使用单个占位符的模板
                    template = "请解释{}的原理和使用场景。"
                    question = template.format(keywords[0] if keywords else np.random.choice(self.tech_keywords[selected_techs[0]]))
            else:
                # 单个占位符
                if keywords:
                    question = template.format(np.random.choice(keywords))
                else:
                    question = template.format(np.random.choice(self.tech_keywords[selected_techs[0]]))
            
            questions.append(question)
            labels.append(list(selected_techs))
        
        # 创建数据集
        dataset = {
            "questions": questions,
            "labels": labels
        }
        
        return dataset
    
    def generate_and_save_dataset(self, num_samples=5000, test_size=0.2):
        """
        生成并保存数据集
        
        Args:
            num_samples: 样本数量，默认5000
            test_size: 测试集比例，默认0.2
        """
        print("生成面试题数据集...")
        dataset = self.generate_interview_questions(num_samples)
        
        # 划分训练集和测试集
        questions = dataset["questions"]
        labels = dataset["labels"]
        
        X_train, X_test, y_train, y_test = train_test_split(
            questions, labels, test_size=test_size, random_state=42
        )
        
        # 保存为JSON格式
        train_data = {"questions": X_train, "labels": y_train}
        test_data = {"questions": X_test, "labels": y_test}
        
        # 保存训练数据集 - 用于模型训练的面试题和对应技术栈标签
        with open(os.path.join(self.output_dir, "train_data.json"), "w", encoding="utf-8") as f:
            json.dump(train_data, f, ensure_ascii=False, indent=2)
        
        # 保存测试数据集 - 用于评估模型性能的面试题和对应技术栈标签
        with open(os.path.join(self.output_dir, "test_data.json"), "w", encoding="utf-8") as f:
            json.dump(test_data, f, ensure_ascii=False, indent=2)
        
        # 保存技术栈列表 - 所有支持的技术栈类别
        with open(os.path.join(self.output_dir, "tech_stacks.json"), "w", encoding="utf-8") as f:
            json.dump(self.tech_stacks, f, ensure_ascii=False, indent=2)
        
        print(f"数据集已保存到 {self.output_dir}")
        print(f"训练集大小: {len(X_train)}")
        print(f"测试集大小: {len(X_test)}")

# 如果直接运行此文件，则生成数据集
if __name__ == "__main__":
    generator = TechStackDatasetGenerator()
    generator.generate_and_save_dataset(num_samples=10000) 